# 使用torchvision库中的MNIST数据集
import torch
import torchvision.datasets as datasets
import torchvision.transforms as transforms
import torch.utils.data as data_utils
import CNN

train_data = datasets.MNIST(root='mnist', train=True, download=True, transform=transforms.ToTensor())
#  root: 数据集存放路径 train: 是否为训练集  download: 是否下载  transform: 数据预处理
test_data = datasets.MNIST(root='mnist', train=False, download=True, transform=transforms.ToTensor())
# print(train_data,  test_data)
# 分批加载
train_loader = data_utils.DataLoader(dataset=train_data, batch_size=64, shuffle=True) # 数据集
test_loader = data_utils.DataLoader(dataset=test_data, batch_size=64, shuffle=True) # 测试集
print(train_loader, test_loader)
cnn = CNN.CNN()
# 使用显卡加速
cnn = cnn.cuda()
# 损失函数
loss_func = torch.nn.CrossEntropyLoss()
# 优化器
optimizer = torch.optim.Adam(cnn.parameters(), lr=0.01)
for epoch in range(5):
    for i, (images, labels) in enumerate(train_loader):
        labels = labels.cuda()
        # 前向传播
        outputs = cnn(images.cuda())
        # 传入输出层节点和真实标签来计算损失函数
        loss = loss_func(outputs, labels)
        # 清空梯度
        optimizer.zero_grad()
        # 反向传播
        loss.backward()
        # 更新参数
        optimizer.step()
    loss_test= 0
    for index2, (images, labels) in enumerate(test_loader):
        labels = labels.cuda()
        outputs = cnn(images.cuda())
        loss_test += loss_func(outputs, labels).item()
    print('第{}轮，测试集的平均损失为：{:.3f}'.format(epoch + 1, loss_test / len(test_loader)))

torch.save(cnn, 'model/cnn.pkl')
